#include "ResourceProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "LocalScene.h"
#include "Game/GameState.h"

void ResourceProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);

	mCurrMask = 0;
	mPrevMask = 0;

	mMinExtents = glm::vec3(FLT_MAX);
	mMaxExtents = glm::vec3(FLT_MIN);

	mOpaqueEntitys = mGameScene->GetEntitysByTypes<LocalEntity, Model, OpaqueObject>();

	mRoughnessMultiplier = 1.f;

	for (uint32 i = 0; i < HALTON_SAMPLES; i++)
		mJitterSamples.push_back(glm::vec2((2.0f * HaltonSequence(2, i) - 1.0f), (2.0f * HaltonSequence(3, i) - 1.0f)));

	uint32 jitterIndex = GGameState->GetFrameAmount() % HALTON_SAMPLES;
	glm::vec2 halton = mJitterSamples[jitterIndex];
	mCurrentJitter = glm::vec2(halton.x / float(GRenderState->GetWidth()), halton.y / float(GRenderState->GetHeight()));
	mPrevJitter = mCurrentJitter;

	mPerFrameUbo.viewInverse = glm::inverse(mGameScene->GetCamera().GetView());
	mPerFrameUbo.projInverse = glm::inverse(mGameScene->GetCamera().GetProj());
	mPerFrameUbo.projView = mGameScene->GetCamera().GetProj() * mGameScene->GetCamera().GetView();
	mPerFrameUbo.projViewInverse = glm::inverse(mPerFrameUbo.projView);
	mPerFrameUbo.prevProjView = mPerFrameUbo.projView;
	mPerFrameUbo.cameraPosition = glm::vec4(mGameScene->GetCamera().GetPosition(), 1.f);
	mPerFrameUbo.currentPrevJitter = glm::vec4(mCurrentJitter, mPrevJitter);

	mPerFrameUbo.light.SetRadius(mLocalScene->mLightRadius);
	mPerFrameUbo.light.SetColor(mLocalScene->mLightColor);
	mPerFrameUbo.light.SetIntensity(mLocalScene->mLightIntensity);
	mPerFrameUbo.light.SetType(mLocalScene->mLightType);
	mPerFrameUbo.light.SetDirection(-mLocalScene->mLightDirection);
	mPerFrameUbo.light.SetPosition(mLocalScene->mLightPosition);
	mPerFrameUbo.light.SetCosThetaInner(glm::cos(glm::radians(mLocalScene->mLightConeAngleInner)));
	mPerFrameUbo.light.SetCosThetaOuter(glm::cos(glm::radians(mLocalScene->mLightConeAngleOuter)));

	mGBufferPushConstants.resize(mOpaqueEntitys.size(), {});

	for (uint32 i = 0; i < mOpaqueEntitys.size(); ++i)
	{
		mGBufferPushConstants[i].objectIndex = i;
		mGBufferPushConstants[i].roughnessMultiplier = mRoughnessMultiplier;
		mGBufferPushConstants[i].prevModel = mOpaqueEntitys[i]->GetTransform().GetWorldMatrix();
		mGBufferPushConstants[i].model = mOpaqueEntitys[i]->GetTransform().GetWorldMatrix();


		glm::mat4 modelMatrix = mOpaqueEntitys[i]->GetTransform().GetWorldMatrix();
		Model& model = mOpaqueEntitys[i]->GetComponent<Model>();
		for (uint32 j = 0; j < model.mVertexs.size(); ++j)
		{
			glm::vec3 tempPos = glm::vec3(modelMatrix * glm::vec4(model.mVertexs[j].Position, 1.f));

			if (tempPos.x > mMaxExtents.x)
				mMaxExtents.x = tempPos.x;
			if (tempPos.y > mMaxExtents.y)
				mMaxExtents.y = tempPos.y;
			if (tempPos.z > mMaxExtents.z)
				mMaxExtents.z = tempPos.z;

			if (tempPos.x < mMinExtents.x)
				mMinExtents.x = tempPos.x;
			if (tempPos.y < mMinExtents.y)
				mMinExtents.y = tempPos.y;
			if (tempPos.z < mMinExtents.z)
				mMinExtents.z = tempPos.z;
		}
	}

	PerFrameUbo perFrameUbo = mPerFrameUbo;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, perFrameUbo](VkCommandBuffer commandBuffer)
		{
			TopLevelAccelerationStructure::Builder topLevelAccelerationStructureBuilder;

			for (uint32 i = 0; i < mOpaqueEntitys.size(); ++i)
			{
				Model& model = mOpaqueEntitys[i]->GetComponent<Model>();

				std::vector<VertexGpu> tempVertexs;
				tempVertexs.resize(model.mVertexs.size());
				for (uint32 j = 0; j < model.mVertexs.size(); ++j)
				{
					tempVertexs[j].position = model.mVertexs[j].Position;
					tempVertexs[j].normal = model.mVertexs[j].Normal;
					tempVertexs[j].tangent = model.mVertexs[j].Tangent;
					tempVertexs[j].uv = model.mVertexs[j].TexC;
				}

				mOpaqueEntitys[i]->mVertexBuffer = VertexBuffer::Builder()
					.SetData(tempVertexs.data())
					.SetSize(tempVertexs.size() * sizeof(VertexGpu))
					.SetRayUse()
					.Build();

				mOpaqueEntitys[i]->mIndexBuffer = IndexBuffer::Builder()
					.SetData(model.mIndices.data())
					.SetSize(model.mIndices.size() * sizeof(uint32))
					.SetIndexCount(model.mIndices.size())
					.SetRayUse()
					.Build();

				ObjectDescGpu rayObjDesc;
				rayObjDesc.vertexAddress = mOpaqueEntitys[i]->mVertexBuffer->GetDeviceAddress();
				rayObjDesc.indexAddress = mOpaqueEntitys[i]->mIndexBuffer->GetDeviceAddress();

				mObjectDescs.push_back(rayObjDesc);

				mMaterials.push_back(mOpaqueEntitys[i]->mMaterial);


				RefCountPtr<AccelerationStructureGeometry> alsGeometry = AccelerationStructureGeometry::Builder()
					.SetVertexFormat(VK_FORMAT_R32G32B32_SFLOAT)
					.SetVertexStride(sizeof(VertexGpu))
					.SetMaxVertex(model.mVertexs.size())
					.SetVertexData(mOpaqueEntitys[i]->mVertexBuffer)
					.SetIndexData(mOpaqueEntitys[i]->mIndexBuffer)
					.Build();

				RefCountPtr<BottomLevelAccelerationStructure> botLevelAS = BottomLevelAccelerationStructure::Builder()
					.AddGeometry(alsGeometry)
					.SetInstanceCustomIndex(i)
					.SetTransform(mOpaqueEntitys[i]->GetTransform().GetWorldMatrix())
					.Build();

				topLevelAccelerationStructureBuilder.AddBlas(botLevelAS);
			}

			mMaterialBuffer = CustomBuffer::Builder()
				.SetData(mMaterials.data())
				.SetSize(sizeof(MaterialGpu) * mMaterials.size())
				.SetUsage(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
				.Build(this, "materials");

			mObjectDescBuffer = CustomBuffer::Builder()
				.SetData(mObjectDescs.data())
				.SetSize(sizeof(ObjectDescGpu) * mObjectDescs.size())
				.SetUsage(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT)
				.Build(this, "objectDescs");

			mTopLevelAS = topLevelAccelerationStructureBuilder
				.Build(this, "topLevelAS");

			PerFrameUbo perFrameUboData = perFrameUbo;

			mPerFrameUboBuffer = CustomBuffer::Builder()
				.SetData(&perFrameUboData)
				.SetSize(sizeof(PerFrameUbo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build(this, "perFrameUbo");


			mTrilinearSampler = RHISampler::Builder()
				.SetMipmapMode(VK_SAMPLER_MIPMAP_MODE_LINEAR)
				.Build(this, "trilinearSampler");

			mBiLinearSampler = RHISampler::Builder()
				.SetMipmapMode(VK_SAMPLER_MIPMAP_MODE_NEAREST)
				.Build(this, "bilinearSampler");

			mNearestSampler = RHISampler::Builder()
				.SetMagFilter(VK_FILTER_NEAREST)
				.SetMinFilter(VK_FILTER_NEAREST)
				.Build(this, "nearestSampler");

			mSobolSequence = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\sobol_256_4d.png")
				.Build(this, "sobolSequence");

			mScramblingRankingSPP1 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_1spp.png")
				.Build(this, "scramblingRankingSPP1");

			mScramblingRankingSPP2 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_2spp.png")
				.Build(this, "scramblingRankingSPP2");

			mScramblingRankingSPP4 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_4spp.png")
				.Build(this, "scramblingRankingSPP4");

			mScramblingRankingSPP8 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_8spp.png")
				.Build(this, "scramblingRankingSPP8");

			mScramblingRankingSPP16 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_16spp.png")
				.Build(this, "scramblingRankingSPP16");

			mScramblingRankingSPP32 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_32spp.png")
				.Build(this, "scramblingRankingSPP32");

			mScramblingRankingSPP64 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_64spp.png")
				.Build(this, "scramblingRankingSPP64");

			mScramblingRankingSPP128 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_128spp.png")
				.Build(this, "scramblingRankingSPP128");

			mScramblingRankingSPP256 = Picture2DTexture::Builder()
				.SetPath("BlueNoise\\scrambling_ranking_128x128_2d_256spp.png")
				.Build(this, "scramblingRankingSPP256");

			mCubemapImage = CubemapTexture::Builder()
				.SetAutoMip()
				.SetPaths({
				"Hdr\\TreeSky_nz.hdr",
				"Hdr\\TreeSky_pz.hdr",
				"Hdr\\TreeSky_py.hdr",
				"Hdr\\TreeSky_ny.hdr",
				"Hdr\\TreeSky_px.hdr",
				"Hdr\\TreeSky_nx.hdr"
					})
				.SetFormat(VK_FORMAT_R32G32B32A32_SFLOAT)
				.Build(this, "cubemapImage");

			mCubemapSampler = RHISampler::Builder()
				.SetAddressModeU(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeV(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeW(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetCompareOp(VK_COMPARE_OP_NEVER)
				.SetMaxLod(mCubemapImage->GetMipLevels())
				.SetBorderColor(VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE)
				.Build(this, "cubemapSampler");
		});
}

void ResourceProcess::OnTick(float deltaTime)
{
	mPrevJitter = mCurrentJitter;
	uint32 jitterIndex = GGameState->GetFrameAmount() % HALTON_SAMPLES;
	glm::vec2 halton = mJitterSamples[jitterIndex];
	mCurrentJitter = glm::vec2(halton.x / float(GRenderState->GetWidth()), halton.y / float(GRenderState->GetHeight()));

	mPerFrameUbo.viewInverse = glm::inverse(mGameScene->GetCamera().GetView());
	mPerFrameUbo.projInverse = glm::inverse(mGameScene->GetCamera().GetProj());
	mPerFrameUbo.prevProjView = mPerFrameUbo.projView;
	mPerFrameUbo.projView = mGameScene->GetCamera().GetProj() * mGameScene->GetCamera().GetView();
	mPerFrameUbo.projViewInverse = glm::inverse(mPerFrameUbo.projView);
	mPerFrameUbo.cameraPosition = glm::vec4(mGameScene->GetCamera().GetPosition(), 1.f);
	mPerFrameUbo.currentPrevJitter = glm::vec4(mCurrentJitter, mPrevJitter);

	mPerFrameUbo.light.SetRadius(mLocalScene->mLightRadius);
	mPerFrameUbo.light.SetColor(mLocalScene->mLightColor);
	mPerFrameUbo.light.SetIntensity(mLocalScene->mLightIntensity);
	mPerFrameUbo.light.SetType(mLocalScene->mLightType);
	mPerFrameUbo.light.SetDirection(-mLocalScene->mLightDirection);
	mPerFrameUbo.light.SetPosition(mLocalScene->mLightPosition);
	mPerFrameUbo.light.SetCosThetaInner(glm::cos(glm::radians(mLocalScene->mLightConeAngleInner)));
	mPerFrameUbo.light.SetCosThetaOuter(glm::cos(glm::radians(mLocalScene->mLightConeAngleOuter)));

	for (uint32 i = 0; i < mOpaqueEntitys.size(); ++i)
	{
		mGBufferPushConstants[i].roughnessMultiplier = mRoughnessMultiplier;
		mGBufferPushConstants[i].prevModel = mGBufferPushConstants[i].model;
		mGBufferPushConstants[i].model = mOpaqueEntitys[i]->GetTransform().GetWorldMatrix();
	}

	PerFrameUbo perFrameUbo = mPerFrameUbo;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, perFrameUbo](VkCommandBuffer commandBuffer)
		{
			PerFrameUbo perFrameUboData = perFrameUbo;

			mPerFrameUboBuffer->WriteMemory(&perFrameUboData);

			mPrevMask = mCurrMask;
			mCurrMask = GRenderState->GetFrameAmount() % 2;

		});

}

void ResourceProcess::OnResize()
{

}
